from datetime import datetime
from openai import OpenAI
import time

class LLMClient:
    def __init__(self):
        self.api_base, self.model, self.api_key = self.get_model_config()
        self.client = OpenAI(api_key=self.api_key, base_url=self.api_base)
        self.system_prompt = self.get_system_prompt()
        self.user_prompt = self.get_user_prompt()

    def get_model_config(self):
        with open("config/大模型配置.txt", "r", encoding='utf-8') as f:
            api_base = f.readline().split("=")[1].strip()
            model = f.readline().split("=")[1].strip()
        with open("config/API_KEY.txt", "r", encoding='utf-8') as f:
            api_key = f.read().strip()
        return api_base, model, api_key

    def get_system_prompt(self):
        return "你需要针对用户想要筛选的推文类型，仔细筛选符合要求的推文，并按照指定格式输出"

    def get_user_prompt(self):
        with open("config/提示词.txt", "r", encoding='utf-8') as f:
            result = f.read()
        tail = f'''
        \n如果没有符合条件的，则回复None。
        今天的日期是{datetime.now().strftime("%Y-%m-%d")}。
        每一行输出格式：
        标题, 公众号, 发布日期, 阅读数, 相关程度
        其中，相关程度只能是1或2或3，用于表示推文符合要求类型的程度，1表示可能相关，2表示比较相关，3表示非常相关。发布日期就用本来的形式，不用换算成精确日期。
        '''
        return result + tail

    def get_llm_response(self, msg="", stream=False, silence=True):
        if silence:
            stream = False
        question = msg + "\n\n" + self.user_prompt
        completion = self.client.chat.completions.create(
            model=self.model,
            stream=stream,  # 流式：true，非流式：false
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": question},
            ],
            temperature=0,  # Lower temperature for higher accuracy and strictness
            max_tokens=1500,  # Adjust as needed
            top_p=None,  # Use nucleus sampling for better results
            frequency_penalty=0.5,  # Penalize frequent tokens to reduce repetition
            presence_penalty=0.5  # Penalize new topic introduction
        )
        if stream:
            for message in completion:
                if message.choices:
                    answer = message.choices[0].delta.content if hasattr(message.choices[0].delta, 'content') else ""
                    for char in answer:
                        print(char, end="", flush=True)
            print()  # Ensure the output ends with a newline
        else:
            answer = completion.choices[0].message.content
            print(answer) if not silence else None
        return answer

if __name__ == "__main__":
    # Example usage
    llm_client = LLMClient()
    print("llm ready")
    tic = time.time()
    # with open(r'F:\Mycode2025\wechat-llm\release\dist\数据\[03月11日13时16分]订阅号消息.txt', 'r', encoding='utf-8') as f:
    #     print(llm_client.get_llm_response(f.read()))
    msg = '''
    上海徐汇
      查看头像
      上海徐汇
      关注徐汇，感知徐汇，徐汇因你而精彩。
        16766篇原创内容
        22个朋友关注
      视频号 : 上海徐汇
        取消关注
        发消息
    消息
    视频号
      #“城”就美好 53
      #汇营商 惠万企 10
      #大城养老徐汇范式 12
      #城市逆生长 9
      #家庭建设 12
        今天
          上海已实现！这笔钱直接发给个人→
          阅读 1.1万  赞 9  
          名额秒光！市民夜校也可以考证了？来看看都有啥课程
          阅读 1834  赞 2  
          办实事传温暖，徐汇区救助站“暖在社区”
          阅读 126  
          “洪水猛兽”？学习工具？26组家庭在游戏中找到亲子沟通新密码！
          阅读 121  赞 6  
        今天
          全国首创！武康大楼电表箱内多了根防火“乾坤圈”
          阅读 2206  赞 57  
          错峰祭扫、文明祭扫，2025年清明祭扫提示请查收
          阅读 487  赞 2  
        今天
          上海首家！快来徐汇滨江体验炫酷AI黑科技，就在模速空间！
          阅读 3211  赞 105  4个朋友分享
          热闹开场！多彩活动绘就“甜灵”春景
          阅读 322  赞 14  
          吃药期间，这种水果最好别碰！
          阅读 586  赞 1  
        今天
          这家医院心理科正式开诊
          阅读 2880  赞 31  
          乐山绿地上演“公益交响曲”
          阅读 382  赞 2  
          针对这个问题，微信发布紧急公告！
          阅读 258  赞 3  
          春天上高速为何更容易犯困？春游出行必看→
          阅读 88  赞 1  
        今天
          具身智能上升为国家战略，养老机器人未来五年会有突破吗？
          阅读 1586  赞 41  
          气温冲上20℃！雨水接踵而至！
          阅读 1022  赞 4  
          免费开放！来这个展看“巾帼”如何逐梦
          阅读 179  赞 1  
          《六姊妹》中老汤家的“糖尿病梦魇”是怎么回事？专家告诉你→
          阅读 236  
    '''
    print(llm_client.get_llm_response(msg, stream=True, silence=False))
    toc = time.time()
    print(f"Time taken: {toc-tic} seconds")